import pandas as pd
import torch
import matplotlib.pyplot as plt


def read_file_set_date(filename) -> pd.DataFrame:
    # 读取文件
    if 'test.csv' in filename:
        df = pd.read_csv(filename, low_memory=False)
    else:
        df = pd.read_excel(filename)

    # 将Date或date列转为datetime类型
    if 'Date' in df.columns:
        df['Date'] = pd.to_datetime(df['Date'])
        df.set_index('Date', inplace=True)
    elif 'date' in df.columns:
        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)
    else:
        raise ValueError('Date column not found')

    return df


def seq_normalization(X: torch.FloatTensor) -> torch.FloatTensor:
    """
    对序列数据做归一化，输入矩阵可以是任意维度，对最后两维在倒数第二维上做标准化
    :param X: 输入矩阵
    :return: 归一化后的矩阵
    """
    from itertools import product
    if X.dim() < 2:
        raise ValueError("输入矩阵的维度必须大于等于2。")

    def normalize_submatrix(submatrix):
        submatrix = torch.cat([(submatrix[:, :4] - torch.mean(submatrix[:, :4])) / (torch.std(submatrix[:, :4]) + 1e-8),
                               (submatrix[:, 4:] - torch.mean(submatrix[:, 4:], dim=0, keepdim=True)) / (
                                       torch.std(submatrix[:, 4:], dim=0, keepdim=True) + 1e-8)], dim=1)
        return submatrix

    if X.dim() == 2:
        X_normalized = normalize_submatrix(X)
    else:
        shape = X.shape[:-2]
        X_normalized = X.clone()
        for idx in product(*[range(dim) for dim in shape]):
            X_normalized[idx] = normalize_submatrix(X[idx])

    return X_normalized


def stock_adjust(df: pd.DataFrame) -> pd.DataFrame:
    """
    通过收盘价后复权价计算开盘最高最低的后复权价格
    :param df:
    :return:
    """
    # 计算复权因子
    df['factor'] = df['adjust_price'] / df['close']

    # 应用复权因子
    df['open_adjust'] = df['open'] * df['factor']
    df['high_adjust'] = df['high'] * df['factor']
    df['low_adjust'] = df['low'] * df['factor']
    df['close_adjust'] = df['close'] * df['factor']

    # 删除原来的价格列
    df = df.drop(columns=['open', 'high', 'low', 'close'])

    # 重命名后复权价格列
    df = df.rename(columns={
        'open_adjust': 'open',
        'high_adjust': 'high',
        'low_adjust': 'low',
        'close_adjust': 'close'
    })
    return df


class LinePlot:
    def __init__(self, title="figure", save_path="plot.png"):
        self.fig, (self.ax1, self.ax2, self.ax3) = plt.subplots(nrows=3, sharex='all')
        self.fig.suptitle(title)
        self.save_path = save_path

    def draw(self, x, base, total_wealth, action, reward, real_reward, stock_reward, close=True, save=True):
        self.ax1.plot(x, base)
        self.ax1.plot(x, total_wealth)
        self.ax1.legend(["Actual", "Model"])
        self.ax1.set_title('Profit')

        self.ax2.plot(x, action)
        self.ax2.set_title('Action')
        self.ax3.plot(x, reward)
        self.ax3.plot(x, real_reward)
        self.ax3.plot(x, stock_reward)
        self.ax3.legend(["With offset", "Real Reward", "Stock Reward"])
        self.ax3.set_title('Env Reward')
        if save:
            with open(self.save_path, 'wb') as f:
                self.fig.savefig(f, format='png')
        if close:
            plt.close(self.fig)


def average_holding_time(position_list):
    holding_time = 0
    total_trades = 0

    for i in range(1, len(position_list)):
        total_trades += position_list[i - 1] - position_list[i] if position_list[i] < position_list[i - 1] else 0
        holding_time += position_list[i]
    return holding_time / max(total_trades, 1)


def init_weights(model):
    from torch.nn import init
    for param in model.parameters():
        if param.dim() > 1:
            init.xavier_uniform_(param)
    return
